import argparse  # 导入参数库
import datetime
import os
import subprocess
import re  # 导入正则表达式库
import queue  # 导入队列库
import socket
import threading  # 导入线程库
from tqdm import tqdm  # 导入进度条库
import prettytable  # 导入表格库
import pyfiglet  # 导入字体库
from colorama import Fore  # 导入颜色库

# 锁用于控制进度条的更新
progress_lock = threading.Lock()
version = "3.2.1"  # 版本号
threads = []  # 线程列表
port_queue = queue.Queue()  # 端口队列
results_queue = queue.Queue()  # 结果队列
# 队列深度拷贝不能有Lock 要不然不能完成深度拷贝
# 深度拷贝方法table_queue =  copy.deepcopy(results_queue)
table = prettytable.PrettyTable(["index","Port","Service_Name","Time"])  # 用于表格输出
# 打印横幅
def banner():
    print(
        Fore.BLUE + f"PortScan {version}   Author: 细节狂魔-精益求精——Youze"
    )
    test = pyfiglet.figlet_format("PortScan", font="calvin_s")
    print(Fore.BLUE + test)


# 命令参数获取
def get_args():
    parser = argparse.ArgumentParser(description="超级端口扫描工具")
    parser.add_argument(
        "-u", "--target", help="目标IP/URL地址", required=True
    )  # required=True表示必须输入
    parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"%(prog)s {version}",
        help="显示版本信息",
    )
    parser.add_argument(
        "-p",
        "--port",
        help="扫描端口范围 [a] all或 [n] normal或 start-end 或 port1,port2,port3",
        default="n",
    )
    parser.add_argument("-o", "--output", help="输出文件", default="")
    parser.add_argument("-n", "--number", help="线程数量", default=1, type=int)
    return parser.parse_args()


# 端口格式检查函数
def port_check(port):
    if port == "n":  # 默认端口
        default_ports = [
            20,
            21,
            22,
            23,
            25,
            53,
            80,
            110,
            143,
            443,
            445,
            993,
            995,
            1723,
            3306,
            3389,
            5432,
            8080,
            8443,
            1521,
            5900,
            6379,
            8888,
            10000,
            33000,
        ]
        return default_ports
    elif port == "a":  # 全端口扫描
        return list(range(1, 65536))  # 65535 + 1
    elif "-" in port:
        start_port, end_port = map(int, port.split("-"))
        if start_port < 1 or start_port > 65535 or end_port < 1 or end_port > 65535:
            print(Fore.RED + "端口范围错误")
            exit()
        elif start_port > end_port:
            start_port, end_port = end_port, start_port  # 交换
        return list(range(start_port, end_port + 1))
    else:
        return [int(p) for p in port.split(",")]


# 扫描端口函数
def scan_port(host, port_queue, progress_bar):
    while not port_queue.empty():  # 队列不为空时循环
        port = port_queue.get()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # sock.settimeout(5)  # 设置超时
        try:
            result = sock.connect_ex((host, port))  # 尝试连接
            time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间
            if result == 0:
                try:
                    # 尝试获取服务名
                    service_name = socket.getservbyport(port)
                except:
                    service_name = "Unknown"
                results_queue.put((host, port,service_name,time))  # 将结果放入结果队列
        except Exception as e:
            print(Fore.RED + f"错误: {e}")
        finally:
            sock.close()  # 确保 socket 被关闭
        with progress_lock:  # 更新进度条时加锁
            progress_bar.update(1)  # 更新进度条

# 多线程扫描函数
def thread_scan(srgs_host, srgs_port, srgs_number):
    # 将srgs_port中的端口放入队列
    for p in srgs_port:
        port_queue.put(p)

    # 使用 tqdm 创建进度条
    with tqdm(total=len(srgs_port), desc="[+]扫描进度", unit="port") as progress_bar:
        # 创建线程
        for i in range(srgs_number):  # 线程数量
            t = threading.Thread(
                target=scan_port, args=(srgs_host, port_queue, progress_bar)  # 传入参数
            )
            threads.append(t)  # 将线程加入线程列表
            t.start()
        # 等待所有线程完成
        for t in threads:
            t.join()

# 写入日志函数
def wirte_log(srgs_output):  # 写入日志
    log_directory = "./PortScan3_log/"
    if not os.path.exists(log_directory):  # 创建日志目录
        os.makedirs(log_directory)   # 创建目录

    # 确定输出文件名
    output_file = srgs_output if srgs_output else srgs_host  # 未指定输出文件名则使用目标地址作为文件名
    file_path = os.path.join(log_directory, f"{output_file}.csv")  # 确定输出文件路径
    try:
        with open(file_path, "w") as f:
            f.write("index,IP:Port,Service_Name,Time\n")  # 写入表头
            length = results_queue.qsize()  # 队列长度
            while not results_queue.empty():
               index = length - results_queue.qsize() + 1  # 序号
               host, port , Service_Name,time = results_queue.get()  # 取出结果
               table.add_row([index,f"{host}:{port}",Service_Name,time])  # 加入表格
               f.write(f"{index},{host}:{port},{Service_Name},{time}\n")  # 写入日志
            print(Fore.GREEN + f"[+]写入日志：{file_path}")
    except IOError as e:
        print(Fore.RED + f"[-]写入日志失败：{e}")  # 错误处理

# 目标地址检查函数
def url_check(srgs_host):
    if "http://" in srgs_host or "https://" in srgs_host:
        # 正则匹配 协议头和域名，不匹配端口和路径
        pattern = r"^(https?://)?([^/:]+)"
        # 匹配目标地址
        ip = re.findall(pattern, srgs_host)[0]
        srgs_host = ip[0] + ip[1]  # 取出域名
        if re.match(pattern, srgs_host):
            # 发送请求获取 IP 地址
            try:
                ip = socket.gethostbyname(srgs_host.split("//")[1])  # 发送请求获取 IP 地址
                print(Fore.GREEN + f"[+]目标地址：{srgs_host} ({ip})")  # 输出目标地址
                return ip
            except:
                print(Fore.RED + f"[-]目标不可达：{srgs_host}")
                exit()

    else:   # 目标地址为 IP 地址
        pattern = r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$"
        if re.match(pattern, srgs_host):
            return srgs_host
        else:
            print(Fore.RED + f"[-]IP 格式错误：{srgs_host}")
            exit()


if __name__ == "__main__":
    banner()  # 打印横幅
    srgs = get_args()  # 获取命令参数
    srgs_host = url_check(srgs.target)  # 目标地址检查
    srgs_port = port_check(srgs.port)    # 端口格式检查
    srgs_output = srgs.output    # 输出文件名
    srgs_number = srgs.number  # 线程数量
    table.align = "l"  # 表格左对齐
    print(Fore.GREEN + "[+]开始扫描")
    try:
        thread_scan(srgs_host, srgs_port, srgs_number)  # 调用多线程函数
    except KeyboardInterrupt:
        print(Fore.RED + "[-]用户中断")
        exit()
    wirte_log(srgs_output)  # 调用写入日志函数
    print(table)
    print(Fore.GREEN + "[+]扫描完成")
